import backtrader as bt
import backtrader.indicators as btind  # 导入策略分析模块
import pandas as pd
import datetime
import akshare as ak

"""
【信号指标取值与多空信号对应关系】
    signal 指标取值>0 → 对应多头 long 信号；
    signal 指标取值<0 → 对应空头 short 信号；
    signal 指标取值=0 → 不发指令；
    
bt.SIGNAL_LONG：
    多头信号用于做多，空头信号用于平仓 close；
    如果系统中同时存在 LONGEXIT 信号类型，SIGNAL_LONG 中的空头信号将不起作用，将会使用 LONGEXIT 中的空头信号来平仓多头，如上面的多条交易信号的例子。
    
【平仓类】确定平仓信号，在下达平仓指令时，优先级高于上面开仓类中的信号
    bt.SIGNAL_LONGEXIT：接收空头信号平仓多头；
    bt.SIGNAL_SHORTEXIT：接收多头信号平仓空头；
"""


def get_data(code):
    df = ak.stock_zh_a_daily(symbol=code, adjust='qfq')
    # df.date = pd.to_datetime(df.date, format='%Y-%m-%d', utc=True)
    # df.set_index('date', inplace=True)
    df.index = pd.to_datetime(df.date)
    df.sort_index(inplace=True)

    return df


# 定义交易信号1
class SMACloseSignal(bt.Indicator):
    lines = ('signal',)  # 声明 signal 线，交易信号放在 signal line 上
    params = (('period', 30),)

    def __init__(self):
        self.lines.signal = self.data - bt.indicators.SMA(period=self.p.period)


# 定义交易信号2
class SMAExitSignal(bt.Indicator):
    lines = ('signal',)
    params = (('p1', 5), ('p2', 30),)

    def __init__(self):
        sma1 = bt.indicators.SMA(period=self.p.p1)
        sma2 = bt.indicators.SMA(period=self.p.p2)
        self.lines.signal = sma1 - sma2


codes = ['sh600276', 'sh600519', 'sh603288']
# 恒瑞医药
data = get_data(codes[0])

cerebro = bt.Cerebro()
st_date = datetime.datetime(2023, 3, 1)
end_date = datetime.datetime(2023, 5, 29)
datafeed = bt.feeds.PandasData(dataname=data, fromdate=st_date, todate=end_date)
cerebro.adddata(datafeed, name='600276.SH')

cerebro.add_signal(bt.SIGNAL_LONG, SMACloseSignal)
cerebro.add_signal(bt.SIGNAL_LONGEXIT, SMAExitSignal)
cerebro.run()
cerebro.plot()
